#!/usr/bin/python
'''
Mulitprocess version of client_stub
'''
import socket
import struct
import select
import os
import sys
import logging
#we need just one connection for ssh, all other connection is combined to ssh tunnel
#only support NO AUTHENTICATION(X'00')
#only support tcp
#only support CONNECT(X'01') CMD

PORT=1080
PIDFILE='/tmp/ht.pid'
LOGFILE='/tmp/ht.log'

reloadcfgflg=None
quitflg=None
clients=None

class ConnectClosed(Exception):pass

log=logging.info

def read(sk, cnt):
    buf=''
    while len(buf)<cnt:
        tp = sk.recv(cnt-len(buf))
        if not tp:
            raise ConnectClosed
        buf += tp
    return buf

def quitht():
    global clients
    for c in clients:
        try:
            os.kill(c, 9)
        except:
            pass
    log("quit...")
    sys.exit(0)

def reloadcfg():
    log("configuration reloaded")

def run(port):
    global reloadcfgflg
    global quitflg
    global clients
    sk = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sk.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sk.bind(('0.0.0.0', port))
#    sk.bind(('127.0.0.1', port))
    sk.listen(10)
    log('Listener pid: %d listening on port %d' % (os.getpid(), port))
    clients = []
    while True:
        try:
            cs, client = sk.accept()
        except IOError,e:
            if e.errno==4:
                if quitflg:
                    quitht()
                    quitflg=False
                if reloadcfgflg:
                    reloadcfg()
                    reloadcfgflg=False
                continue
            else:
                raise e
        except Exception,e:
            log('Listener Error while accept: %s' % e)
            sk.shutdown(socket.SHUT_RDWR)
            sk.close()
            return 1
        log('client %s:%d connected.' % client)
        pid = os.fork()
        if pid:
            # parent process
            log('client %s:%d child process %d' % (client[0], client[1], pid))
            clients.append(pid)
            cs.close()
            cs = None
        else:
            # child process
            clients = None
            try:
                sk.close()
                # recv ver and method count
                buf = read(cs, 2)
                ver, mtd_cnt = struct.unpack('!BB', buf[:2])
                log('client %s:%d use protocl version: %d' % (client[0], client[1], ver))
                if ver != 5:
                    log('client %s:%d protocol version unsupported.' % client)
                    cs.close()
                    return
                buf = read(cs, mtd_cnt)
                log('client %s:%d supported auth method: %s' % (client[0], client[1], repr(buf)))
                for mtd in buf:
                    (mtd,) = struct.unpack('!B', mtd)
                    if mtd == 0:
                        cs.send('\x05\x00')
                        break
                else:
                    log('client %s:%d all auth method unsupported.' % client)
                    cs.send('\x05\xFF')
                    cs.close()
                    return
                buf = read(cs, 4)
                ver, cmd, rsv, atyp = struct.unpack('!BBBB', buf)
                log('client %s:%d CMD:%d, atyp:%d' % (client[0], client[1], cmd, atyp))
                if cmd != 1:
                    log('client %s:%d unsupported CMD:%d' % (client[1], client[1], cmd))
                    cs.send('\x05\x07\x00\x01\xFF\xFF\xFF\xFF\x00\x00')
                    cs.close()
                    return
                if atyp == 1:
                    buf = read(cs, 4+2)
                    addr = struct.unpack('!IH', buf)
                    addr = (str(addr[0]),addr[1])
                elif atyp == 3:
                    buf = read(cs, 1)
                    (addrlen,) = struct.unpack('!B', buf)
                    log('client %s:%d connect dst addr lenght:%d' % (client[0], client[1], addrlen))
                    buf = read(cs, addrlen+2)
                    addr = struct.unpack('!%dsH' % addrlen, buf)
                else:
                    log('client %s:%d unsupported address type:%d' % (client[0], client[1], atyp))
                    cs.send('\x05\x08\x00\x01\xFF\xFF\xFF\xFF\x00\x00')
                    cs.close()
                    return

                log('client %s:%d request connect to %s:%d' % (client[0], client[1], addr[0], addr[1]))
                #connect to addr
                ss = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                ss.connect(addr)
                log('client %s:%d connected to %s:%d' % (client[0], client[1], addr[0], addr[1]))
                # Send some http traffic to cheat firewall
                '''ss.send('\r\n'.join(['GET / HTTP/1.1',
                                     'Host: %s' % addr[0],
                                     'Connection: keep-alive',
                                     'User-Agent: Mozilla/5.0 (X11; U; Linux i686; en-US) AppleWebKit/532.2 (KHTML, like Gecko) Chrome/4.0.221.8 Safari/532.2',
                                     'Accept: application/xml,application/xhtml+xml,text/html;q=0.9,text/plain;q=0.8,image/png,*/*;q=0.5',
                                     'Accept-Encoding: gzip,deflate',
                                     'Accept-Language: zh-CN,zh;q=0.8',
                                     'Accept-Charset: gb18030,utf-8;q=0.7,*;q=0.3',
                                     ])
                        )
                ss.send('\r\n')
                # wait for server response
                '''
                
                cs.send('\x05\x00\x00\x01\xFF\xFF\xFF\xFF\x00\x00')
                #exchange data
                while True:
                    [rs,],ws,xs = select.select([cs,ss],[],[])
                    if rs is cs:
                        ws = ss
                    else:
                        ws = cs
                    buf = rs.recv(4096)
                    if not buf:
                        ss.close()
                        cs.close()
                        raise ConnectClosed
                    ws.send(buf)
            except ConnectClosed,e:
                log('client %s:%d closed' % client)
                return 0
            except Exception,e:
                log('client %s:%d exception catched: %s' % (client[0], client[1], str(e)))
                return 1
def daemonize():
    if os.fork():
        exit(0)
    os.setsid()
    #fork again
    if os.fork():
        exit(0)
    os.chdir('/')
    os.umask(0)
    sys.stdin.close()
    sys.stdout.close()
    sys.stderr.close()
    #catch some signal
    import signal
    def sigint_func(x, frame):
        global reloadcfgflg
        log("sigint catched")
        reloadcfgflg=True
    signal.signal(signal.SIGINT, sigint_func)
    def sigterm_func(x, frame):
        global quitflg
        log("sigterm catched")
        quitflg=True
    signal.signal(signal.SIGTERM, sigterm_func)
    signal.signal(signal.SIGCHLD, signal.SIG_IGN)
    
def main():
    daemonize()
    logging.basicConfig(level=logging.DEBUG,
                        format='%(asctime)s %(levelname)s %(message)s',
                        filename=LOGFILE)
    open(PIDFILE,'w').write('%d' % os.getpid())
    run(PORT)

if __name__ == '__main__':
    main()

